import os
import json
from .map import label_list
def extract_data(gt_dir,pred_dir):
    names = os.listdir(gt_dir)

    # 用于存放金标准数据的列表
    gt_boxes = []
    gt_labels = []
    difficulties = []
    gt_nodule_density = [] # 此处为平均密度
    gt_group = []
    gt_layer = []
    gt_volume = []
    gt_totalVolume = []

    # 用于存放预测结果的列表
    ai_boxes = []
    ai_labels = []
    ai_scores = []
    ai_nodule_density = [] # 此处为平均密度
    norpr = []
    ai_group = []
    ai_layer = []
    ai_volume = []
    ai_totalVolume = []



    # 导入json中的数据
    gt_list = []
    pred_list = []
    for name in names:
        pred_path = os.path.join(pred_dir,name)
        gt_path = os.path.join(gt_dir,name)
        with open(gt_path,'r') as f:
            gt = json.load(f)
        with open(pred_path,'r') as f:
            pred = json.load(f)
        gt_list += [gt]
        pred_list += [pred]
    
    for i in range(len(gt_list)):
        # 金标准
        gt_boxes.append([gt_list[i]["box"][x] for x in range(len(gt_list[i]["label"])) if gt_list[i]["label"][x] in label_list[:-1]])
        gt_labels.append([gt_list[i]["label"][x] for x in range(len(gt_list[i]["label"])) if gt_list[i]["label"][x] in label_list[:-1]])
        difficulties.append([gt_list[i]["difficult"][x] for x in range(len(gt_list[i]["label"])) if gt_list[i]["label"][x] in label_list[:-1]])
        gt_nodule_density.append([gt_list[i]["nodule_density"][x] for x in range(len(gt_list[i]["label"])) if gt_list[i]["label"][x] in label_list[:-1]]) 
        gt_group.append([int(gt_list[i]["group"][x]) for x in range(len(gt_list[i]["label"])) if gt_list[i]["label"][x] in label_list[:-1]])
        gt_layer.append([gt_list[i]["layer"][x] for x in range(len(gt_list[i]["label"])) if gt_list[i]["label"][x] in label_list[:-1]])
        gt_volume.append([gt_list[i]["volume"][x] for x in range(len(gt_list[i]["label"])) if gt_list[i]["label"][x] in label_list[:-1]])
        gt_totalVolume.append([gt_list[i]["totalVolume"][x] for x in range(len(gt_list[i]["label"])) if gt_list[i]["label"][x] in label_list[:-1]])

        # 预测结果
        ai_boxes.append([pred_list[i]["box"][x] for x in range(len(pred_list[i]["label"])) if pred_list[i]["label"][x] in label_list[:-1]])
        ai_labels.append([pred_list[i]["label"][x] for x in range(len(pred_list[i]["label"])) if pred_list[i]["label"][x] in label_list[:-1]])
        ai_scores.append([pred_list[i]["score"][x] for x in range(len(pred_list[i]["label"])) if pred_list[i]["label"][x] in label_list[:-1]])
        ai_nodule_density.append([pred_list[i]["nodule_density"][x] for x in range(len(pred_list[i]["label"])) if pred_list[i]["label"][x] in label_list[:-1]])
        norpr.append([pred_list[i]["norpr"][x] for x in range(len(pred_list[i]["label"])) if pred_list[i]["label"][x] in label_list[:-1]])
        ai_group.append([int(pred_list[i]["group"][x]) for x in range(len(pred_list[i]["label"])) if pred_list[i]["label"][x] in label_list[:-1]])
        ai_layer.append([pred_list[i]["layer"][x] for x in range(len(pred_list[i]["label"])) if pred_list[i]["label"][x] in label_list[:-1]])
        ai_volume.append([pred_list[i]["volume"][x] for x in range(len(pred_list[i]["label"])) if pred_list[i]["label"][x] in label_list[:-1]])
        ai_totalVolume.append([pred_list[i]["totalVolume"][x] for x in range(len(pred_list[i]["label"])) if pred_list[i]["label"][x] in label_list[:-1]])
        
        gt_data = {
            "split":'gt',
            "ct_name":names,
            "gt_boxes":gt_boxes,# 手工边界、bbox的两种标注可能同时存在，按照肺结节测试方案v3.1.pptx手工边界优先级高于bbox
            "gt_labels":gt_labels,
            "difficulties":difficulties,
            "gt_nodule_density":gt_nodule_density,
            "gt_group":gt_group,
            "gt_layer":gt_layer,
            "gt_volume":gt_volume,
            "gt_totalVolume":gt_totalVolume,
        }
        ai_data = {
            "split":'ai',
            "ct_name":names,
            "ai_boxes":ai_boxes,
            "ai_labels":ai_labels,
            "ai_scores":ai_scores,
            "ai_nodule_density":ai_nodule_density,
            "ai_norpr":norpr,
            "ai_group":ai_group,
            "ai_layer":ai_layer,
            "ai_volume":ai_volume,
            "ai_totalVolume":ai_totalVolume,
        }
    return gt_data,ai_data,names

if __name__ == '__main__':
    gt_dir = 'data/new_gt_1657004830051'
    pred_dir = 'data/new_ai_1657004830051'
    gt_data,ai_data,names = extract_data(gt_dir,pred_dir)